"""知识图谱模型定义"""

from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Float, JSON
from sqlalchemy.orm import relationship
from sqlalchemy.dialects.postgresql import UUID
from typing import Optional, Dict, Any, List
import uuid
from datetime import datetime
from src.config.database import Base
from sqlalchemy import Index

class KnowledgeGraphNode(Base):
    """知识图谱节点模型"""
    __tablename__ = "knowledge_graph_nodes"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    content = Column(Text, nullable=False)
    content_type = Column(String(50), default="text")
    embedding = Column(JSON)  # 存储嵌入向量
    node_metadata = Column(JSON)  # 存储节点元数据
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 添加索引以提高查询性能
    __table_args__ = (
        Index('idx_kg_node_content_type', 'content_type'),
        Index('idx_kg_node_created_at', 'created_at'),
        Index('idx_kg_node_metadata_gin', 'node_metadata', postgresql_using='gin'),
    )

    # 关系
    outgoing_edges = relationship("KnowledgeGraphEdge", 
                                foreign_keys="KnowledgeGraphEdge.source_node_id",
                                back_populates="source_node")
    incoming_edges = relationship("KnowledgeGraphEdge", 
                                foreign_keys="KnowledgeGraphEdge.target_node_id",
                                back_populates="target_node")


class KnowledgeGraphEdge(Base):
    """知识图谱边模型"""
    __tablename__ = "knowledge_graph_edges"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    source_node_id = Column(UUID(as_uuid=True), ForeignKey("knowledge_graph_nodes.id"), nullable=False)
    target_node_id = Column(UUID(as_uuid=True), ForeignKey("knowledge_graph_nodes.id"), nullable=False)
    relationship_type = Column(String(100), nullable=False)
    edge_metadata = Column(JSON)  # 存储边元数据
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关系
    source_node = relationship("KnowledgeGraphNode", 
                              foreign_keys=[source_node_id],
                              back_populates="outgoing_edges")
    target_node = relationship("KnowledgeGraphNode", 
                              foreign_keys=[target_node_id],
                              back_populates="incoming_edges")